/**
 * Copyright 2019-2020 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "dwconv2d_quant_process_fusion_pass.h"
#include "bias_optimize_quant_rollback_base.h"

namespace fe {
/*
 *  fusion pattern
 *           input
 *               \
 *             AscendQuant
 *                \
 *                 v
 *     weights--->DepthwiseConv2D--->AscendDequant--->output
 *                ^                    ^
 *               /                    /
 *              /                    /
 *           bias                deq_scale
 *
 *  fusion pattern1
 *           input
 *               \
 *             AscendQuant
 *                \
 *                Pad
 *                  \
 *                   v
 *     weights--->DepthwiseConv2D--->AscendDequant--->output
 *                 /                       ^
 *                /                       /
 *               /                       /
 *             bias                  deq_scale
 */
vector<FusionPattern *> DWConv2DQuantProcessFusionPass::DefinePatterns() {
  vector<FusionPattern *> patterns;
  FE_LOGD("Start to define DepthwiseConv2D quant process fusion pattern.");
  FusionPattern *pattern = new (std::nothrow) FusionPattern("DWConv2DQuantProcessFusion");
  FE_CHECK(pattern == nullptr, FE_LOGW("new FusionPattern object failed!"), return patterns);
  pattern->AddOpDesc(PATTERN_QUANT, {QUANT})
      .AddOpDesc(PATTERN_CUBE, {DEPTHWISECONV2D})
      .AddOpDesc(PATTERN_DEQUANT, {DEQUANT})
      .SetInputs(PATTERN_CUBE, {PATTERN_QUANT})
      .SetInputs(PATTERN_DEQUANT, {PATTERN_CUBE})
      .SetOutput(PATTERN_DEQUANT);
  patterns.push_back(pattern);

  FusionPattern *pattern1 = new (std::nothrow) FusionPattern("DWConv2DQuantProcessFusion1");
  FE_CHECK(pattern1 == nullptr, FE_LOGW("new FusionPattern object failed!"), return patterns);
  pattern1->AddOpDesc(PATTERN_QUANT, {QUANT})
      .AddOpDesc(PATTERN_PAD, {PAD})
      .AddOpDesc(PATTERN_CUBE, {DEPTHWISECONV2D})
      .AddOpDesc(PATTERN_DEQUANT, {DEQUANT})
      .SetInputs(PATTERN_PAD, {PATTERN_QUANT})
      .SetInputs(PATTERN_CUBE, {PATTERN_PAD})
      .SetInputs(PATTERN_DEQUANT, {PATTERN_CUBE})
      .SetOutput(PATTERN_DEQUANT);
  patterns.push_back(pattern1);

  return patterns;
}

/*
 * DepthwiseConv2D may have Pad node before
 * so we need to refresh Pad node dtype
 */
Status DWConv2DQuantProcessFusionPass::SetPadNodeDataType(ge::NodePtr &pad_node) {
  ge::OpDescPtr pad_op_desc = pad_node->GetOpDesc();
  ge::GeTensorDesc pad_input_desc = pad_op_desc->GetInputDesc(0);
  ge::GeTensorDesc pad_output_desc = pad_op_desc->GetOutputDesc(0);
  pad_input_desc.SetDataType(ge::DT_FLOAT16);
  pad_input_desc.SetOriginDataType(ge::DT_FLOAT16);
  pad_output_desc.SetDataType(ge::DT_FLOAT16);
  pad_output_desc.SetOriginDataType(ge::DT_FLOAT16);
  if (pad_op_desc->UpdateInputDesc(0, pad_input_desc) != SUCCESS) {
    REPORT_FE_ERROR("[GraphOpt][DWConv2dQntPcsFus][SetPadNdDtype] Fail to update input desc of Node[%s].",
                    pad_op_desc->GetName().c_str());
    return FAILED;
  }
  if (pad_op_desc->UpdateOutputDesc(0, pad_output_desc) != SUCCESS) {
    REPORT_FE_ERROR("[GraphOpt][DWConv2dQntPcsFus][SetPadNdDtype] Fail to update output desc of Node[%s].",
                    pad_op_desc->GetName().c_str());
    return FAILED;
  }
  FE_LOGD("Data type of Pad[%s] has been updated to float during quant rollback", pad_op_desc->GetName().c_str());
  return SUCCESS;
}

Status DWConv2DQuantProcessFusionPass::SetDataTypeOfNodes(ge::NodePtr &cube_node) {
  ge::DataType data_type = ge::DT_FLOAT;
  ge::DataType target_data_type = ge::DT_FLOAT16;

  if (SetCubeNodeDataType(cube_node, data_type, target_data_type) != SUCCESS) {
    REPORT_FE_ERROR("[GraphOpt][DWConv2dQntPcsFus][SetDtype] Set cube node[%s] data type failed.",
                    cube_node->GetName().c_str());
    return FAILED;
  }
  FE_LOGD("Set cube node[%s] data type success.", cube_node->GetName().c_str());

  ge::NodePtr cube_input_node = cube_node->GetInAllNodes().at(0);
  if (cube_input_node->GetType() == PAD) {
    if (SetPadNodeDataType(cube_input_node) != SUCCESS) {
      REPORT_FE_ERROR("[GraphOpt][DWConv2dQntPcsFus][SetDtype] Set node[%s] data type failed.",
                      cube_input_node->GetName().c_str());
      return FAILED;
    }
    FE_LOGD("Set pad node[%s] data type success.", cube_input_node->GetName().c_str());
  }
  return SUCCESS;
}

ge::NodePtr DWConv2DQuantProcessFusionPass::GetCubeNodeInputNode(ge::NodePtr &cube_node) {
  ge::NodePtr node_after_quant = cube_node;
  ge::NodePtr cube_input_node = cube_node->GetInAllNodes().at(0);
  if (cube_input_node->GetType() == PAD) {
    node_after_quant = cube_input_node;
  }
  return node_after_quant;
}

void DWConv2DQuantProcessFusionPass::SetCinCoutReverse(ge::NodePtr &nodePtr) {
  (void)ge::AttrUtils::SetBool(nodePtr->GetOpDesc(), ATTR_CIN_COUT_REVERSE, true);
}

}  // namespace fe